#!/usr/bin/env python3
"""
下载开源时间序列数据集并自动划分训练测试集
"""

import os
import numpy as np
import pandas as pd
import requests
import zipfile
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
import matplotlib.pyplot as plt
import seaborn as sns

# 设置matplotlib中文字体
plt.rcParams['font.sans-serif'] = ['SimHei', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False

def download_ucr_dataset(dataset_name='ECG200'):
    """
    下载UCR时间序列分类数据集
    
    Args:
        dataset_name (str): 数据集名称，默认为ECG200（心电图数据）
    
    Returns:
        tuple: (X_train, y_train, X_test, y_test)
    """
    print(f"正在下载 {dataset_name} 数据集...")
    
    # UCR数据集URL
    base_url = "http://www.timeseriesclassification.com/Downloads/"
    zip_url = f"{base_url}{dataset_name}.zip"
    
    # 创建数据目录
    data_dir = f"data/{dataset_name}"
    os.makedirs(data_dir, exist_ok=True)
    
    try:
        # 下载zip文件
        response = requests.get(zip_url, stream=True)
        zip_path = os.path.join(data_dir, f"{dataset_name}.zip")
        
        with open(zip_path, 'wb') as f:
            for chunk in response.iter_content(chunk_size=8192):
                f.write(chunk)
        
        # 解压文件
        with zipfile.ZipFile(zip_path, 'r') as zip_ref:
            zip_ref.extractall(data_dir)
        
        print(f"数据集下载完成: {dataset_name}")
        
        # 读取训练和测试数据
        train_file = os.path.join(data_dir, f"{dataset_name}_TRAIN.txt")
        test_file = os.path.join(data_dir, f"{dataset_name}_TEST.txt")
        
        if os.path.exists(train_file) and os.path.exists(test_file):
            # 读取训练数据
            train_data = pd.read_csv(train_file, header=None, delim_whitespace=True)
            X_train = train_data.iloc[:, 1:].values  # 特征
            y_train = train_data.iloc[:, 0].values   # 标签
            
            # 读取测试数据
            test_data = pd.read_csv(test_file, header=None, delim_whitespace=True)
            X_test = test_data.iloc[:, 1:].values    # 特征
            y_test = test_data.iloc[:, 0].values     # 标签
            
            print(f"训练集大小: {X_train.shape}")
            print(f"测试集大小: {X_test.shape}")
            print(f"类别分布 - 训练集: {np.unique(y_train, return_counts=True)}")
            print(f"类别分布 - 测试集: {np.unique(y_test, return_counts=True)}")
            
            return X_train, y_train, X_test, y_test
        else:
            print("数据集文件不存在，尝试生成模拟数据")
            return generate_synthetic_data()
            
    except Exception as e:
        print(f"下载数据集失败: {e}")
        print("使用模拟数据替代")
        return generate_synthetic_data()

def generate_synthetic_data():
    """
    生成模拟时间序列数据
    
    Returns:
        tuple: (X_train, y_train, X_test, y_test)
    """
    print("生成模拟时间序列数据...")
    
    np.random.seed(42)
    n_samples = 1000
    sequence_length = 50
    n_features = 5
    
    # 生成特征数据
    X = np.random.randn(n_samples, sequence_length, n_features)
    
    # 生成标签（二分类问题）
    # 基于时间序列的模式生成标签
    y = np.zeros(n_samples)
    for i in range(n_samples):
        # 如果序列的平均值大于0，则标记为1类
        if np.mean(X[i, :, 0]) > 0:
            y[i] = 1
        # 如果序列的方差较大，则标记为1类
        elif np.var(X[i, :, 1]) > 1.0:
            y[i] = 1
    
    # 划分训练测试集
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.2, random_state=42, stratify=y
    )
    
    print(f"模拟数据生成完成:")
    print(f"训练集大小: {X_train.shape}")
    print(f"测试集大小: {X_test.shape}")
    print(f"类别分布 - 训练集: {np.unique(y_train, return_counts=True)}")
    print(f"类别分布 - 测试集: {np.unique(y_test, return_counts=True)}")
    
    return X_train, y_train, X_test, y_test

def save_dataset_to_csv(X_train, y_train, X_test, y_test, dataset_name='synthetic'):
    """
    将数据集保存为CSV文件
    
    Args:
        X_train: 训练特征
        y_train: 训练标签
        X_test: 测试特征
        y_test: 测试标签
        dataset_name: 数据集名称
    """
    # 保存训练集
    train_df = pd.DataFrame(X_train.reshape(X_train.shape[0], -1))
    train_df['target'] = y_train
    train_file = f"data/{dataset_name}_train.csv"
    train_df.to_csv(train_file, index=False)
    
    # 保存测试集
    test_df = pd.DataFrame(X_test.reshape(X_test.shape[0], -1))
    test_df['target'] = y_test
    test_file = f"data/{dataset_name}_test.csv"
    test_df.to_csv(test_file, index=False)
    
    print(f"训练集已保存: {train_file}")
    print(f"测试集已保存: {test_file}")
    
    # 也保存一个合并的样本数据文件
    sample_df = pd.concat([train_df.head(100), test_df.head(50)])
    sample_file = "data/sample_data.csv"
    sample_df.to_csv(sample_file, index=False)
    print(f"样本数据已保存: {sample_file}")

def visualize_dataset(X_train, y_train, X_test, y_test, dataset_name='synthetic'):
    """
    可视化数据集
    
    Args:
        X_train: 训练特征
        y_train: 训练标签
        X_test: 测试特征
        y_test: 测试标签
        dataset_name: 数据集名称
    """
    print("生成数据可视化图表...")
    
    # 创建可视化目录
    os.makedirs("data/visualization", exist_ok=True)
    
    # 1. 类别分布图
    plt.figure(figsize=(12, 4))
    
    plt.subplot(1, 3, 1)
    train_counts = pd.Series(y_train).value_counts()
    plt.bar(train_counts.index, train_counts.values)
    plt.title('Train Set Class Distribution')
    plt.xlabel('Class')
    plt.ylabel('Count')
    
    plt.subplot(1, 3, 2)
    test_counts = pd.Series(y_test).value_counts()
    plt.bar(test_counts.index, test_counts.values)
    plt.title('Test Set Class Distribution')
    plt.xlabel('Class')
    plt.ylabel('Count')
    
    plt.subplot(1, 3, 3)
    total_counts = pd.Series(np.concatenate([y_train, y_test])).value_counts()
    plt.bar(total_counts.index, total_counts.values)
    plt.title('Overall Class Distribution')
    plt.xlabel('Class')
    plt.ylabel('Count')
    
    plt.tight_layout()
    plt.savefig('data/visualization/class_distribution.png', dpi=300, bbox_inches='tight')
    plt.close()
    
    # 2. 时间序列示例图
    plt.figure(figsize=(12, 8))
    
    # 从每个类别中选取几个样本进行可视化
    unique_classes = np.unique(y_train)
    for i, cls in enumerate(unique_classes):
        # 找到该类别的样本索引
        cls_indices = np.where(y_train == cls)[0]
        
        # 随机选择3个样本
        sample_indices = np.random.choice(cls_indices, min(3, len(cls_indices)), replace=False)
        
        for j, idx in enumerate(sample_indices):
            plt.subplot(len(unique_classes), 3, i * 3 + j + 1)
            
            # 绘制第一个特征的时间序列
            plt.plot(X_train[idx, :, 0])
            plt.title(f'Class {cls} - Sample {j+1}')
            plt.xlabel('Time Step')
            plt.ylabel('Value')
    
    plt.tight_layout()
    plt.savefig('data/visualization/time_series_examples.png', dpi=300, bbox_inches='tight')
    plt.close()
    
    print("数据可视化完成")

def main():
    """主函数"""
    print("=" * 60)
    print("开源时间序列数据集下载和划分工具")
    print("=" * 60)
    
    # 尝试下载UCR数据集
    try:
        X_train, y_train, X_test, y_test = download_ucr_dataset('ECG200')
        dataset_name = 'ECG200'
    except:
        # 如果下载失败，使用模拟数据
        X_train, y_train, X_test, y_test = generate_synthetic_data()
        dataset_name = 'synthetic'
    
    # 保存数据集
    save_dataset_to_csv(X_train, y_train, X_test, y_test, dataset_name)
    
    # 可视化数据集
    visualize_dataset(X_train, y_train, X_test, y_test, dataset_name)
    
    # 生成数据统计报告
    print("\n" + "=" * 60)
    print("数据统计报告:")
    print("=" * 60)
    print(f"数据集名称: {dataset_name}")
    print(f"训练集大小: {X_train.shape}")
    print(f"测试集大小: {X_test.shape}")
    print(f"序列长度: {X_train.shape[1]}")
    print(f"特征数量: {X_train.shape[2]}")
    print(f"类别数量: {len(np.unique(np.concatenate([y_train, y_test])))}")
    print(f"训练集类别分布: {dict(zip(*np.unique(y_train, return_counts=True)))}")
    print(f"测试集类别分布: {dict(zip(*np.unique(y_test, return_counts=True)))}")
    
    print("\n数据集准备完成！")
    print("生成的文件:")
    print("- data/ECG200_train.csv (训练集)")
    print("- data/ECG200_test.csv (测试集)")
    print("- data/sample_data.csv (样本数据)")
    print("- data/visualization/class_distribution.png (类别分布图)")
    print("- data/visualization/time_series_examples.png (时间序列示例图)")

if __name__ == "__main__":
    main()